import numpy as np
from environment.env import theta_0

# _pi 为agent的策略, 模块内的全局变量
_pi=None

def simple_convert_into_pi_from_theta(theta):
    '''theta_0 转换为 pi'''
    global _pi
    [m, n] = theta.shape  # 
    _pi = np.zeros((m, n))
    for i in range(0, m):
        _pi[i, :] = theta[i, :] / np.nansum(theta[i, :])  # 

    _pi = np.nan_to_num(_pi)  # nan


# _pi 为agent的策略
# _pi= 
simple_convert_into_pi_from_theta(theta_0)




def get_next_s(s):
    # 策略pi 为模块内的全局变量
    direction = ["up", "right", "down", "left"]

    next_direction = np.random.choice(direction, p=_pi[s, :])
    # _pi[s,:]

    if next_direction == "up":
        s_next = s - 3  # 上
    elif next_direction == "right":
        s_next = s + 1  # 右
    elif next_direction == "down":
        s_next = s + 3  # 下
    elif next_direction == "left":
        s_next = s - 1  # 左

    return s_next








